# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/13 17:17 
@Author  : ZhangShenao 
@File    : 1.使用Pipeline调用模型.py 
@Desc    : 使用高层接口——Pipeline调用模型
"""
import dotenv
import torch
from transformers import pipeline

# 加载环境变量
dotenv.load_dotenv()

# 设置cuda|cpu模式
device = "cuda" if torch.cuda.is_available() else "cpu"

# 构造一个Pipeline
# 用途是text-generation文本生成
# 引用的模型是Qwen/Qwen2.5-0.5B-Instruct
messages = [
    {"role": "user", "content": "请写一首赞美秋天的五言绝句"},
]
pipe = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct", device=device, max_new_tokens=100)

# 使用Pipeline调用模型,获取结果
result = pipe(messages)
print(result[-1]['generated_text'][-1]['content'])
